
from channels.generic.websocket import AsyncWebsocketConsumer
from celery.result import AsyncResult
from asgiref.sync import sync_to_async
import json,pickle

def get_task_info(task_id):
    try:
        task = AsyncResult(task_id)
        status = task.status
        result = task.result if task.ready() else None
        return task_id, status, result

    except Exception as e:
        print(f"Error getting task info for {task_id}: {e}")
        raise


class TaskStatusConsumer(AsyncWebsocketConsumer):
    # authentication_classes = []
    async def connect(self):
        await self.accept()

    async def disconnect(self, close_code):
        pass

    async def receive(self, text_data):
        # print("type:",type(text_data), text_data)
        # 假设从前端接收的是任务ID列表，这里简化处理
        # print(text_data)
        task_ids = text_data.split(',')
        # task_ids = list(text_data)
        # print("type:",type(task_ids))
        # print("task_ids:", task_ids)
        for task_id in task_ids:
            # print("task_id", task_id)
            try:
                # 使用 sync_to_async 包装 get_task_info 函数
                task_id, status, result = await sync_to_async(get_task_info)(task_id)
                print("task_id", task_id, "status", status, "result", result)
                await self.send(text_data=f"{task_id}:{status}:{result}")
            except Exception as e:
                print(e)